package com.ztom.agent.config;

import com.ztom.agent.util.PkgExpUtils;
import com.ztom.agent.util.StrMatchUtils;
import org.apache.commons.lang3.StringUtils;

import java.util.HashSet;
import java.util.Set;

/**
 * @author ZhangTao
 */
public class TraceFilter {

    /**
     * 不需要注入的 [Package前缀] 集合
     */
    private static Set<String> excludePackagePrefix = new HashSet<>();

    /**
     * 不需要注入的 [Package表达式] 集合
     */
    private static Set<String> excludePackageExp = new HashSet<>();

    /**
     * 需要注入的 [Package前缀] 集合
     */
    private static Set<String> includePackagePrefix = new HashSet<>();

    /**
     * 需要注入的 [Package表达式] 集合
     */
    private static Set<String> includePackageExp = new HashSet<>();

    /**
     * 不需要注入的 [method] 集合
     */
    private static Set<String> excludeMethods = new HashSet<>();

    /**
     * 需要加载的 [plugin] 集合
     */
    private static Set<String> plugins = new HashSet<>();

    /**
     * 不注入的 [ClassLoader] 集合
     */
    private static Set<String> excludeClassLoader = new HashSet<>();

    static {
        // 默认不注入的 package
        excludePackagePrefix.add("java/");
//        excludePackagePrefix.add("javax/");
        excludePackagePrefix.add("sun/");
        excludePackagePrefix.add("com/sun/");
        excludePackagePrefix.add("com/intellij/");

        // 不注入 MyTraceAgent 本身
        excludePackagePrefix.add("com/ztom/agent/");

        // 默认注入的 package
        includePackagePrefix.add("net/paoding/rose/jade/context/JadeInvocationHandler");
        includePackagePrefix.add("org/apache/ibatis/binding/MapperProxy");
        includePackagePrefix.add("com/alibaba/dubbo/rpc/proxy/InvokerInvocationHandler");
        includePackagePrefix.add("org/apache/dubbo/rpc/proxy/InvokerInvocationHandler");
        includePackagePrefix.add("com/alipay/sofa/rpc/proxy/jdk/JDKInvocationHandler");
        includePackagePrefix.add("com/weibo/api/motan/proxy/RefererInvocationHandler");

        // 默认不注入的method
        excludeMethods.add("main");
        excludeMethods.add("premain");
        // java.lang.Object
        excludeMethods.add("getClass");
        excludeMethods.add("hashCode");
        excludeMethods.add("equals");
        excludeMethods.add("clone");
        excludeMethods.add("toString");
        excludeMethods.add("notify");
        excludeMethods.add("notifyAll");
        excludeMethods.add("wait");
        excludeMethods.add("finalize");
        // spring
        excludeMethods.add("afterPropertiesSet");
    }

    /**
     * 插件是否需要加载
     *
     * @param innerClassName
     * @return
     */
    public static boolean isNeedLoad(String innerClassName) {
        if (innerClassName == null) {
            return false;
        }

        return isMatch(innerClassName, plugins);
    }

    private static boolean isMatch(String innerClassName, Set<String> plugins) {
        for (String plugin : plugins) {
            if (StrMatchUtils.isMatch(innerClassName, plugin)) {
                return true;
            }
        }
        return false;
    }

    public static void addPlugins(String plugin) {
        if (StringUtils.isBlank(plugin)) {
            return;
        }

        String packagePrefix = "com.ztom.agent.asm.plugin*";
        Set<String> pkgSet = PkgExpUtils.parse(packagePrefix + plugin);
        plugins.addAll(pkgSet);
    }



    /**
     * 判断该类是否属于不需要扫描的包
     *
     * @param innerClassName : 形如: com/ztom/agent/config/TraceFilter
     * @return : true->不需要修改字节码  false->需要修改字节码
     */
    public static boolean isNotNeedInject(String innerClassName) {
        if (innerClassName == null) {
            return true;
        }

        if (innerClassName.indexOf('$') >= 0) {
            // 内部类 不注入
            return true;
        }

        return isMatch(innerClassName, excludePackagePrefix, excludePackageExp);
    }

    /**
     * 判断该类是否属于需要扫描的包
     *
     * @param innerClassName : 形如: com/ztom/agent/config/TraceFilter
     * @return : true->需要修改字节码  false->不需要修改字节码
     */
    public static boolean isNeedInject(String innerClassName) {
        if (innerClassName == null) {
            return false;
        }

        return isMatch(innerClassName, includePackagePrefix, includePackageExp);
    }

    private static boolean isMatch(String innerClassName, Set<String> pkgPrefixSet, Set<String> pkgExpSet) {
        for (String prefix : pkgPrefixSet) {
            if (innerClassName.startsWith(prefix)) {
                return true;
            }
        }

        for (String exp : pkgExpSet) {
            if (StrMatchUtils.isMatch(innerClassName, exp)) {
                return true;
            }
        }
        return false;
    }

    /**
     * 添加不需要扫描的包
     *
     * @param pkg
     */
    public static void addExcludePackage(String pkg) {
        if (StringUtils.isEmpty(pkg)) {
            return;
        }

        addPackages(pkg, excludePackagePrefix, excludePackageExp);
    }

    private static void addPackages(String packages, Set<String> pkgPrefixSet, Set<String> pkgExpSet) {
        Set<String> pkgSet = PkgExpUtils.parse(packages);
        for (String pkg : pkgSet) {
            pkg = preProcess(pkg);
            if (pkg.indexOf('*') > 0) {
                pkgExpSet.add(pkg);
            } else {
                pkgPrefixSet.add(pkg);
            }
        }
    }

    private static String preProcess(String pkg) {
        return pkg.replace('.', '/').trim();
    }

    /**
     * 添加需要扫描的包
     *
     * @param pkg
     */
    public static void addIncludePackage(String pkg) {
        if (StringUtils.isEmpty(pkg)) {
            return;
        }

        addPackages(pkg, includePackagePrefix, includePackageExp);
    }

    /**
     * 获取不需要扫描的包集合
     *
     * @return
     */
    public static Set<String> getExcludePackagePrefix() {
        return new HashSet<>(excludePackagePrefix);
    }

    /**
     * 获取需要扫描的包前缀集合
     *
     * @return
     */
    public static Set<String> getIncludePackagePrefix() {
        return new HashSet<>(includePackagePrefix);
    }

    /**
     * 获取需要扫描的包表达式集合
     *
     * @return
     */
    public static Set<String> getIncludePackageExp() {
        return new HashSet<>(includePackageExp);
    }



    /**
     * 判断该方法是否属于不需要注入
     *
     * @param methodName
     * @return : true->需要修改字节码  false->不需要修改字节码
     */
    public static boolean isNotNeedInjectMethod(String methodName) {
        if (methodName == null) {
            return false;
        }

        if (methodName.indexOf('$') >= 0) {
            return true;
        }

        return excludeMethods.contains(methodName);
    }

    /**
     * 添加不需要注入的方法
     *
     * @param method
     */
    public static void addExcludeMethods(String method) {
        if (method == null) {
            return;
        }

        excludeMethods.add(method.trim());
    }

    /**
     * 获取不需要注入的方法集合
     *
     * @return
     */
    public static Set<String> getExcludeMethods() {
        return new HashSet<>(excludeMethods);
    }


    /**
     * 添加不需要扫描的类加载器
     *
     * @param classLoader
     */
    public static void addExcludeClassLoader(String classLoader) {
        excludeClassLoader.add(classLoader);
    }


    /**
     * 判断该加载器是否属于不需要注入的类加载器
     *
     * @param classLoader
     * @return : true->需要修改字节码  false->不需要修改字节码
     */
    public static boolean isNotNeedInjectClassLoader(String classLoader) {
        return excludeClassLoader.contains(classLoader);
    }
}
